#!/usr/bin/env python
#coding=utf-8
from __future__ import with_statement
'''
try:
    import psyco
    psyco.full()
except ImportError:
    print 'Psyco not installed, the program will just run slower'
'''
from threading import Lock, Thread
from threading import currentThread
from time import time, sleep
from multiprocessing import Process
from multiprocessing.connection import Listener, Client
import socket
from Queue import Queue
import cPickle as pickle
import os, sys

DEBUG = False


class SessionItem(object):
    liveTime = DEBUG and 10 or 20 * 60
    
    __slots__ = "sessionId", "expire", "_locks", "_lockslock", "_data"
        
    def __init__(self, sessionId):
        self.sessionId = sessionId
        self.expire = time() + self.liveTime
        self._locks = {}
        self._lockslock = Lock()
        self._data = {}
        
    def lock(self, key):
        self.expire = time() + self.liveTime
        while True:
            try:
                self._locks[key][1] += 1
                self._locks[key][0].acquire()
                return
            except KeyError:
                with self._lockslock:
                    if key not in self._locks:
                        self._locks[key] = [Lock(), 0]
                                    
    def unlock(self, key):
        self.expire = time() + self.liveTime
        self._locks[key][0].release()
        self._locks[key][1] -= 1
        if self._locks[key][1] == 0:
            with self._lockslock:
                if self._locks[key][1] == 0:
                    del self._locks[key]
            
    def __contains__(self, key):
        self.expire = time() + self.liveTime
        return key in self._data
        
    def __getitem__(self, key):
        self.expire = time() + self.liveTime
        return self._data[key]
    
    def __setitem__(self, key, value):
        self.expire = time() + self.liveTime
        self._data[key] = value
        
    def __delitem__(self, key):
        self.expire = time() + self.liveTime
        del self._data[key]
        
IN = 0
READ = 1
WRITE = 2
REMOVE = 3
LOCK = 6
UNLOCK = 7

ADDRESS = sys.platform == 'win32' and ('localhost', 8685,) or r'/tmp/simple.session.simple'

        
class SessionServer(Thread):
        def __init__(self, address=ADDRESS, processes=16, daemon=True):
            self._sessions = {}
            self.processes = processes
            self._address = address
            self._stop = False
            super(SessionServer, self).__init__()
            self.daemon = daemon
        '''    
        def __del__(self):
            self.stop()
            sleep(0.1)
            try:
                self.terminate()
            except:
                pass
        '''
        class Worker(Thread):
            def __init__(self, server, conn):
                self.conn = conn
                self._server = server
                self._sessions = self._server._sessions
                super(SessionServer.Worker, self).__init__()
                
            def _proceess(self, sessionId, opcode, key, value=None):
                if sessionId not in self._sessions:
                    self._sessions[sessionId] = SessionItem(sessionId)
                item = self._sessions[sessionId]
                if opcode == IN:
                    return key in item
                elif opcode == READ:
                    return item[key]
                elif opcode == WRITE:
                    item[key] = value
                    return
                elif opcode == REMOVE:
                    del item[key]
                    return
                elif opcode == LOCK:
                    return item.lock(key)
                elif opcode == UNLOCK:
                    return item.unlock(key)            
                
            def run(self):
                conn = self.conn
                while not self._server._stop:
                    try:
                        args = conn.recv()
                        r = None
                        if isinstance(args, list) or isinstance(args, tuple):
                            r = self._proceess(*args)
                        elif isinstance(args, dict):
                            r = self._proceess(**args)
                        else:
                            conn.close()
                            return
                        conn.send((r, None, None,))
                    except Exception, ex:
                        try:
                            if DEBUG:
                                print type(ex), ex
                            conn.send((None, type(ex), ex,))
                        except Exception, ex:
                            if DEBUG:
                                print type(ex), ex
                            try:
                                conn.close()
                                return
                            except Exception, ex:
                                print type(ex), ex
                                return
                conn.close()

        def run(self):
            print 'start to listen session at', self._address
            try:
                os.unlink(self._address)
            except:
                pass
            try:
                self._listener = Listener(self._address)
            except Exception, ex:
                print self._address
                raise ex
            
            self.cleaner = self.Cleaner(self)
            self.cleaner.start()
            while not self._stop:
                try:
                    conn = self._listener.accept()
                    worker = self.Worker(self, conn)
                    worker.start()
                except Exception, ex:
                    try:
                        print type(ex), ex
                        conn.close()
                    except Exception, ex:
                        print type(ex), ex
            self._listener.close()
            return self

        def spawn(self):
            def _proceess():
                self.start()
                self.join()
            self.mainProcess = Process(target=_proceess)
            self.mainProcess.daemon = True
            self.mainProcess.start()
            self.terminate = lambda : self.mainProcess.terminate()
            return self
        
        def config(self, address):
            self._address = address
                
        def stop(self):
            self._stop = True
            
        class Cleaner(Thread):
            sleepTime = DEBUG and 10 or 60
            def __init__(self, sessionServer):
                self.sessionServer = sessionServer
                super(SessionServer.Cleaner, self).__init__()
                
            def run(self):
                while not self.sessionServer._stop:
                    sleep(self.sleepTime)
                    now = time()
                    for (sessionId, sessionItem) in self.sessionServer._sessions.items():
                        if sessionItem.expire < now:
                            del self.sessionServer._sessions[sessionId]

class SessionClient(object):
    def __init__(self, address):
        self._address = address
        self._client = Client(self._address)
        self._lock = Lock()

    def _rpc(self, *args):
        with self._lock:
            try:
                self._client.send(args)
            except:
                try:
                    self._client.close()
                except Exception, ex:
                    print type(ex), ex
                self._client = Client(self._address)
                self._client.send(args)
            (r, Ex, ex,) = self._client.recv()

        if Ex or ex:
            raise Ex, ex
        return r        
            
    def has(self, sessionId, key):
        return self._rpc(sessionId, IN, key)
        
    def get(self, sessionId, key):
        return self._rpc(sessionId, READ, key)
            
    def set(self, sessionId, key, value):
        return self._rpc(sessionId, WRITE, key, value)
        
    def remove(self, sessionId, key):
        return self._rpc(sessionId, REMOVE, key)
            
    def lock(self, sessionId, key):
        return self._rpc(sessionId, LOCK, key)
        
    def unlock(self, sessionId, key):
        return self._rpc(sessionId, UNLOCK, key)

class SessionBase(object):
    def __init__(self, sessionName='SESSION_ID', address=ADDRESS, addresses=[], processMode=sys.platform != 'win32', startServer=True, daemon=True):
        self.sessionName = sessionName
        self._processMode = processMode
        self.address = address
        if addresses:
            self.addresses = addresses        
        else:
            self.addresses = [address,]
        self._startServer = startServer
        self._daemon = daemon
        
    def config(self, addresses, processMode=None, startServer=None, sessionName=None, daemon=None):
        self.addresses = addresses
        if processMode is not None:
            self._processMode = processMode
        if startServer is not None:
            self._startServer = startServer
        if sessionName is not None:
            self.sessionName = sessionName
        if daemon is not None:
            self._daemon = daemon
            
    def start(self):
        if self._startServer:
            address = self.address
            self._server = SessionServer(address)
            if self._processMode:
                self._server.spawn()
            else:
                if self._daemon:
                    self._server.daemon = self._daemon
                self._server.start()
        self._clients = []
        self._size = 0
        for addr in self.addresses:
            self._size += 1
            self._clients.append(SessionClient(addr))
    '''    
    def __del__(self):
        try:
            self.close()
        except:
            pass
    '''    
    def _client(self, sid):
        if self._size == 1:
            return self._clients[0]
        return self._clients[hash(sid) / 7 % self._size]
        
    def close(self):
        self._server.stop()
        sleep(0.1)
        if self._processMode:
            try:
                self._server.terminate()
            except:
                pass

    def getSessionId(self):
        raise NotImplementedError
    
    def __contains__(self, key):
        sid = self.getSessionId()
        return self._client(sid).has(sid, key)
    
    def __getitem__(self, key):
        sid = self.getSessionId()
        return self._client(sid).get(sid, key)
    
    def __setitem__(self, key, value):
        sid = self.getSessionId()
        return self._client(sid).set(sid, key, value)
    
    def __delitem__(self, key):
        sid = self.getSessionId()
        return self._client(sid).remove(sid, key)
    
    def lock(self, key):
        sid = self.getSessionId()
        return self._client(sid).lock(sid, key)
    
    def unlock(self, key):
        sid = self.getSessionId()
        return self._client(sid).unlock(sid, key)
    
def waitForDie():
    while True:
        try:
            raw_input()
        except KeyboardInterrupt:
            return


if __name__ == '__main__':
    session = SessionServer()
    if len(sys.argv) == 2:
        address = sys.argv[1]
        session.config(address)
    if len(sys.argv) == 3:
        address = (sys.argv[1], int(sys.argv[2]))
        session.config(address)
    session.start()
    waitForDie()